#! /usr/bin/env python
#
# Implement cookbook using pyfitsio
import sys
import os

import ctypes

# Use just the C-types bindings to CFITSIO
from fitsio import cfitsio, pyfitsio

def main():
    '''
    Select rows from an input table and copy to the output table
    '''

    infilename  = 'atestfil.fit'
    outfilename = 'btestfil.fit'

    # Open the existing FITS files
    ffin  = pyfitsio.open(infilename, mode='r')
    ffout = pyfitsio.open(outfilename, mode='rw')

    # Move to 3rd HDU in the input file
    hdutype = ffin.move_to_hdu(3)
    if not hdutype == cfitsio.BINARY_TBL:
        print 'Error: expected to find a binary table in this HDU'
        sys.exit(1)
    print ffin.read_axes()

    # Move to the 2nd HDU in output file
    ffout.move_to_hdu(2)

    # Create new extension in the output file
    ffout.create_hdu()

    # Get number of keywords in input file
    status = ctypes.c_int(0)
    nkeys = ctypes.c_int()
    keypos = ctypes.c_int()
    if cfitsio.fits_get_hdrpos(ffin.fptr, ctypes.byref(nkeys), \
                               ctypes.byref(keypos), ctypes.byref(status)):
        pyfitsio.handle_error(ffin, status)

    # Copy all keywords from input to the output extension
    card = ctypes.create_string_buffer('\000' * cfitsio.FLEN_CARD)
    for i in range(1, nkeys.value+1):
        cfitsio.fits_read_record(ffin.fptr, i, card, ctypes.byref(status))
        cfitsio.fits_write_record(ffout.fptr, card, ctypes.byref(status))

    # Read NAXIS1 and NAXIS2 keywords to get table size
    #print ffin.read_axes()
    naxis1, naxis2 = ffin.read_axes()

    # Find which column cotains the density values
    colnum = ctypes.c_int()
    cfitsio.fits_get_colnum(ffin.fptr, cfitsio.CASEINSEN, "density", \
                                ctypes.byref(colnum), \
                                ctypes.byref(status))



    # Read the density column values
    frow = ctypes.c_longlong(1)
    felem = ctypes.c_longlong(1)
    nullval = ctypes.c_float(-99.)
    density = (ctypes.c_float * naxis2)()
    anynulls = ctypes.c_int()
    cfitsio.fits_read_col(ffin.fptr, cfitsio.TFLOAT, colnum, frow, felem, \
                              naxis2, ctypes.byref(nullval), density, \
                              ctypes.byref(anynulls), ctypes.byref(status))


    # allocate buffer large enough for 1 row of the table
    buff = (ctypes.c_ubyte * naxis1)()

    # If the density is less than 3.0, copy the row to the output table
    noutrows = ctypes.c_longlong(0)
    for irow in range(1, naxis2+1):
        if density[irow-1] < 3.0:
            noutrows.value += 1
            cfitsio.fits_read_tblbytes(ffin.fptr, irow, 1, naxis1, buff, \
                                       ctypes.byref(status))
            cfitsio.fits_write_tblbytes(ffout.fptr, noutrows, 1, naxis1, buff, \
                                        ctypes.byref(status))
    
    # Update NAXIS2 keyword with the correct number of rows. Note use of
    # None to set comment argument as NULL pointer
    ffout.write_key('NAXIS2', noutrows.value, None)

    if status.value > 0:
        pyfitsio.handle_error(ffin, status)

    ffout.close()
    ffin.close()

if __name__ == '__main__':
    main()
